import base64
import logging
from Crypto.Cipher import AES
from Crypto.Protocol.KDF import PBKDF2
from Crypto.Random import get_random_bytes
from typing import Optional

logger = logging.getLogger(__name__)

class CryptoUtils:
    """加密工具类，使用AES-256-GCM进行数据加密"""
    
    @staticmethod
    def derive_key(password: str, salt: bytes, key_length: int = 32) -> bytes:
        """
        使用PBKDF2从密码派生密钥
        
        Args:
            password: 密码
            salt: 盐值
            key_length: 密钥长度
            
        Returns:
            bytes: 派生密钥
        """
        return PBKDF2(password, salt, key_length, count=100000)
    
    @staticmethod
    def encrypt_data(data: str, password: str) -> Optional[str]:
        """
        加密数据
        
        Args:
            data: 要加密的数据
            password: 加密密码
            
        Returns:
            Optional[str]: 加密后的数据(Base64编码)
        """
        if not data or not password:
            return None
            
        try:
            # 生成盐和随机数
            salt = get_random_bytes(16)
            nonce = get_random_bytes(16)
            
            # 派生密钥
            key = CryptoUtils.derive_key(password, salt)
            
            # 创建加密器并加密数据
            cipher = AES.new(key, AES.MODE_GCM, nonce=nonce)
            ciphertext, tag = cipher.encrypt_and_digest(data.encode('utf-8'))
            
            # 组合所有部分: salt + nonce + tag + ciphertext
            encrypted_data = salt + nonce + tag + ciphertext
            
            # Base64编码返回
            result = base64.b64encode(encrypted_data).decode('utf-8')
            logger.info("数据加密成功")
            return result
            
        except Exception as e:
            logger.error(f"数据加密失败: {e}")
            return None
    
    @staticmethod
    def decrypt_data(encrypted_data: str, password: str) -> Optional[str]:
        """
        解密数据
        
        Args:
            encrypted_data: 加密的数据(Base64编码)
            password: 解密密码
            
        Returns:
            Optional[str]: 解密后的数据
        """
        if not encrypted_data or not password:
            return None
            
        try:
            # Base64解码
            encrypted_data = base64.b64decode(encrypted_data)
            
            # 提取各部分
            salt = encrypted_data[:16]
            nonce = encrypted_data[16:32]
            tag = encrypted_data[32:48]
            ciphertext = encrypted_data[48:]
            
            # 派生密钥
            key = CryptoUtils.derive_key(password, salt)
            
            # 创建解密器并解密数据
            cipher = AES.new(key, AES.MODE_GCM, nonce=nonce)
            decrypted_data = cipher.decrypt_and_verify(ciphertext, tag)
            
            result = decrypted_data.decode('utf-8')
            logger.info("数据解密成功")
            return result
            
        except Exception as e:
            logger.error(f"数据解密失败: {e}")
            return None